#! /usr/bin/python

# To change this template, choose Tools | Templates
# and open the template in the editor.

__author__="MrEction"
__date__ ="$Apr 12, 2011 7:22:07 PM$"

import pygame
from hexconfig import *
#import datahandler

REPLAY = "04121652"
REPLAY_LIST = ["petesshortreplays/04282359", "petesshortreplays/04290000", "petesshortreplays/04290001", "petesshortreplays/04290002", "petesshortreplays/04290003", "petesshortreplays/0429000538", "petesshortreplays/0429000600", "petesshortreplays/0429000620", "petesshortreplays/0429000636", "petesshortreplays/0429000723", "petesshortreplays/0429000737", "petesshortreplays/0429000823", "petesshortreplays/0429000849", "petesshortreplays/0429000910"]
#REPLAY_LIST.extend(["newreplays/04261738","newreplays/04261742","newreplays/04261759","newreplays/04261805","newreplays/04261822", "newreplays/04261847", "newreplays/04261855", "newreplays/04261906"])
#REPLAY_LIST.extend(["newerreplays/04281222","newerreplays/04282005","newerreplays/04282015","newerreplays/05010453","newerreplays/05010507","newerreplays/05010514","newerreplays/05010523","newerreplays/05010528","newerreplays/05010532"])
#REPLAY_LIST.extend(["goodreplays/04070202","goodreplays/04070205","goodreplays/04070228","goodreplays/04070311","goodreplays/04070543","goodreplays/04121731","goodreplays/04121743","goodreplays/04122038"])
#REPLAY_LIST.extend(["goodreplays/04131831","goodreplays/04132052","goodreplays/04161520","goodreplays/04161930","goodreplays/04161944","goodreplays/04161947","goodreplays/04161957","goodreplays/04252226"])
pygame.init()

SAVE_DATA = True
FLESH_OUT_DATA = False
PROPOGATE_UTILITY = True

PRINT_DATA_SET = False

if __name__ == "__main__":
    print "Train with Replays"
    clock = pygame.time.Clock()
    clock.tick()
    screen = pygame.display.set_mode((SWIDTH,SHEIGHT))#,pygame.FULLSCREEN)#,pygame.HWSURFACE)

    #dh = DataHandler()
    #ds = dh.autoLoadDataSet()#
    loadFile = open("../ai_data/fleshedNodes.pkl",'rb')
    nodes = pickle.load(loadFile)
    print "loaded nodes", len(nodes)
    ds = DataSet()
    ds.nodes = nodes
    keys = ds.nodes.keys()
    for key in keys:
        if key[0]!= "W":
            ds.nodes[key].u = 0.0
            ds.nodes[key].w = 1.0

    
    propogationList = []
    allEndBoardStates = [[] for i in range(2)]
    for r in REPLAY_LIST:
        replay = loadReplay(r)
        
        currentState = replay.setups
        endBoardStates = [[] for i in range(2)]
        turn = 0
        for e in replay.moveList:
            nextState = stateStringFromActionList(currentState,e)
            ds.addNode(nextState)
            #print nextState
            endBoardStates[turn].append(nextState)
            allEndBoardStates[turn].append(nextState)

            if nextState[0] == "W":
                propogationList.append((endBoardStates[turn][:-1], 1.0))
                propogationList.append((endBoardStates[1 if turn == 0 else 0][-3:], -1.0))
                break;

            unitTeam = teamKilledByActionList(currentState,e)
            if unitTeam==0:
                propogationList.append((endBoardStates[turn][-2:], -.5))
                propogationList.append((endBoardStates[1 if turn == 0 else 0][-4:],  .5))
            elif unitTeam==1:
                propogationList.append((endBoardStates[turn][-4:],  .5))
                propogationList.append((endBoardStates[1 if turn == 0 else 0][-2:], -.5))

            currentState = nextState
            currentState = switchStateStringToOtherTeam(currentState, turn)
            turn = (turn+1)%2
        print "end: ", endBoardStates
    ms = clock.tick()
    secs = ms/1000
    mins = secs/60
    print
    print "TIME TO ADD DATA:",mins,"mins,",secs%60 if secs>60 else secs,"secs"
    print "NODES_LENGTH",len(ds.nodes)
    
    #if SAVE_DATA: dh.autoSaveDataSet(ds)

    #ds.fleshOutEdges()
    clock = pygame.time.Clock()
    dateTime = "DS_"+strftime("%m_%d(%H-%M)")
    p = pickle.Pickler(open("../ai_data/backup_datasets/"+dateTime+"Nodes"+".pkl",'wb'))
    p.fast = True
    #p.dump(ds.nodes)

    """
    ms = clock.tick()
    secs = ms/1000
    mins = secs/60
    print
    print "TIME TO FLESH OUT DATA:",mins,"mins,",secs%60 if secs>60 else secs,"secs"
    print "NODES_LENGTH",len(ds.nodes)
    #if SAVE_DATA: dh.autoSaveDataSet(ds)
    """

    if PROPOGATE_UTILITY:
        print "Going to propogate!", len(propogationList)# len(allEndBoardStates)
        clock.tick()
        for prop in propogationList:
            print "I'm about to propogate!"
            ds.propogateUtil(prop[0], prop[1])

        for i in range(len(allEndBoardStates)):
            print "adding util to end states!", i
            for endStateString in allEndBoardStates[i]:
                node = ds.getNode(endStateString)
                node.addUtil(node.util()+0.1,0.1)
        ms = clock.tick()
        secs = ms/1000
        mins = secs/60
        print
        print "TIME TO PROPOGATE UTIL: ",mins,"mins,",secs%60 if secs>60 else secs,"secs"
        print "NODES_LENGTH",len(ds.nodes)

        print "UTIL PROPOGATION PRINTOUTS:"
        numWithUtil = 0
        keys = ds.nodes.keys()
        for k in keys:
            if not ds.nodes[k].util()==0.0:
                numWithUtil+=1
                #print ds.nodes[k].util(), ds.nodes[k].stateString, ds.nodes[k].w
        print "  ", numWithUtil, "/", len(keys), "nodes have utility"



    #if SAVE_DATA: dh.autoSaveDataSet(ds)

    #ds.fleshOutEdges()
    clock = pygame.time.Clock()
    dateTime = "DS_"+strftime("%m_%d(%H-%M)")
    p = pickle.Pickler(open("../ai_data/new0point2Nodes"+".pkl",'wb'))
    p.fast = True
    p.dump(ds.nodes)

